# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-08-29

from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, precision_score, recall_score, \
    classification_report


class Evaluator:
    """
    用于评估分类模型性能的工具类。
    :param y_true: 真实的标签序列。
    :param y_pred: 模型预测的标签序列。
    :param y_prob: 模型预测的标签概率序列，用于计算AUC-ROC（如果需要）。
    """
    def __init__(
            self,
            y_true,
            y_pred,
            y_prob=None,
    ):
        self.y_true = y_true
        self.y_pred = y_pred
        self.y_prob = y_prob  # 概率预测，用于计算AUC-ROC

    def f1(
            self,
            average='binary',
    ):
        """
        计算F1分数。
        :param average: str, 多类分类的平均策略。
        :return: float, 计算出的F1分数。
        """
        return f1_score(self.y_true, self.y_pred, average=average)

    def auc_roc(
            self,
    ):
        """
        计算AUC-ROC分数。
        :return: float, 计算出的AUC-ROC分数。
        :raises ValueError: 如果没有提供预测概率y_prob。
        """
        if self.y_prob is None:
            raise ValueError("y_prob is required for AUC-ROC calculation.")
        return roc_auc_score(self.y_true, self.y_prob)

    def accuracy(
            self,
    ):
        """
        计算准确率。
        :return: float, 计算出的准确率。
        """
        return accuracy_score(self.y_true, self.y_pred)

    def precision(
            self,
            average='binary',
    ):
        """
        计算精确率。
        :param average: str, 多类分类的平均策略。
        :return: float, 计算出的精确率。
        """
        return precision_score(self.y_true, self.y_pred, average=average)

    def recall(
            self,
            average='binary',
    ):
        """
        计算召回率。
        :param average: str, 多类分类的平均策略。
        :return: float, 计算出的召回率。
        """
        return recall_score(self.y_true, self.y_pred, average=average)

    def print_classification_report(
            self,
    ):
        """
        打印分类报告，包括主要的分类指标。
        """
        print(classification_report(self.y_true, self.y_pred))


# 调用示例，假设 y_true 和 y_pred 是已经定义好的标签和预测结果
# y_true = [0, 1, 1, 0, 1]
# y_pred = [0, 1, 0, 0, 1]
# y_prob = [0.1, 0.4, 0.35, 0.8, 0.7]  # 概率预测
# evaluator = Evaluator(y_true, y_pred, y_prob)
# print("F1 Score:", evaluator.f1())
# print("AUC-ROC:", evaluator.auc_roc())
# print("Accuracy:", evaluator.accuracy())
# print("Precision:", evaluator.precision())
# print("Recall:", evaluator.recall())
